package heap_trianning

import (
	"sort"
)

type Customer struct {
	ID    int
	Buy   int
	Time  int
}
type Daddys []*Customer
type Cands  []*Customer

func GetTopKCustomers(customers []int, operate []bool, K int) [][]int {
	customersMap := map[int]*Customer{} // 记录客户是否已经买过东西
	daddys := make(Daddys, 0)           // 得奖区的客户
	cands := make(Cands, 0)             // 候奖区的客户
	ans := make([][]int, 0)             // 初始化结果
	if K == 0 {
		return ans
	}
	for i, c := range customers {
		if _, ok := customersMap[c]; !ok && !operate[i] { // 不存在该客户且退货的，直接跳过
			ans = append(ans, daddys.GetCustomers())
			continue
		}
		// 获取到当前的顾客，没有的话初始化一个顾客
		var curC  *Customer
		if cc, ok := customersMap[c]; ok {
			curC = cc
		} else {
			curC = &Customer{ID: c}
			customersMap[c] = curC
		}
		if operate[i] { // 如果当前是买操作
			curC.Buy++
		} else {
			curC.Buy--
		}
		if curC.Buy == 0 { // 如果此时顾客买的数量为0，情况顾客的map
			delete(customersMap, curC.ID)
		}
		// 将当前顾客不在得奖区与候奖区的情况排除
		if !daddys.Contain(curC) && !cands.Contain(curC) {
			if len(daddys) < K {
				curC.Time = i // 赋值进入得奖区的时间
				daddys = append(daddys, curC)
			} else {
				curC.Time = i
				cands = append(cands, curC)
			}
		}
		daddys.CleanZeroBuy()                    // 清除得奖区里退货变成0购买的
		cands.CleanZeroBuy()                     // 清除候奖区里退货变成0购买的
		sort.Sort(cands)                         // 得奖区重新调整顺序
		sort.Sort(daddys)                        // 得奖区重新调整顺序
		move(&daddys, &cands, K, i)              // 如果候奖区与得奖区需要替换，则需要移动元素
		ans = append(ans, daddys.GetCustomers()) // 重新获取得奖区的顾客
	}
	return ans
}

func move(daddys *Daddys, cands *Cands, K int, time int) {
	if len(*cands) == 0 { // 候选区为空，不需要移动
		return
	}
	if len(*daddys) < K { // 得奖区没满的情况，一定是候选区有，但是得奖区其中有一个顾客退货
		(*cands)[0].Time = time
		*daddys = append(*daddys, (*cands)[0])
		*cands = (*cands)[1:]
	} else { // 得奖区满了，且候奖区有客户的情况，如果候奖区的顾客最大的购买量超过得奖区的顾客，则发生移动
		if (*cands)[0].Buy > (*daddys)[0].Buy { // 侯奖区与得奖区客户发生交换
			oldDaddy := (*daddys)[0]
			newDaddy := (*cands)[0]
			oldDaddy.Time = time
			newDaddy.Time = time
			(*daddys)[0], (*cands)[0] = (*cands)[0], (*daddys)[0]
		}
	}
}

func (d Daddys) Len() int {
	return len(d)
}

func (d Daddys) Less(i, j int) bool {
	if d[i].Buy != d[j].Buy {
		return d[i].Buy < d[j].Buy
	}
	return d[i].Time < d[j].Time
}

func (d Daddys) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

func (d *Daddys) Contain(cc *Customer) bool {
	var ans bool
	for _, dc := range *d {
		if dc == cc {
			ans = true
		}
	}
	return ans
}

func (d *Daddys) CleanZeroBuy()  {
	var newD Daddys
	for _, dc := range *d {
		if dc.Buy != 0 {
			newD = append(newD, dc)
		}
	}
	*d = newD
}

func (d *Daddys) GetCustomers() []int {
	var customers []int
	for _, dc := range *d {
		customers = append(customers, dc.ID)
	}
	return customers
}

func (c Cands) Len() int {
	return len(c)
}

func (c Cands) Less(i, j int) bool {
	if c[i].Buy != c[j].Buy {
		return c[i].Buy > c[j].Buy
	}
	return c[i].Time < c[j].Time
}

func (c Cands) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

func (c *Cands) Contain(cc *Customer) bool {
	var ans bool
	for _, candc := range *c {
		if candc == cc {
			ans = true
		}
	}
	return ans
}

func (c *Cands) CleanZeroBuy()  {
	var newD Cands
	for _, dc := range *c {
		if dc.Buy != 0 {
			newD = append(newD, dc)
		}
	}
	*c = newD
}

func (c *Cands) GetCustomers() []int {
	var customers []int
	for _, dc := range *c {
		customers = append(customers, dc.ID)
	}
	return customers
}

